import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from model import *

# 选择训练设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 准备数据
train_data = torchvision.datasets.CIFAR10("./dataset_2", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10("./dataset_2", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)

print(f"训练数据长度：{len(train_data)}")
print(f"测试数据长度：{len(test_data)}")

# 获取数据
train_loader = DataLoader(train_data, 64)
test_loader = DataLoader(test_data, 64)

# 实例化网络
my_nn = MyNN()
my_nn = my_nn.to(device)
# 损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)
# 优化器
optimizer = torch.optim.SGD(my_nn.parameters(), lr=0.01)

# 训练轮数
epoch = 10
# 记录训练次数
train_step = 0
# 记录测试次数
test_step = 0
# 记录日志
writer = SummaryWriter("logs_train")

# 训练流程
for i in range(epoch):
    print(f"--------------第{i + 1}轮训练开始-----------------")

    for imgs, targets in train_loader:
        imgs = imgs.to(device)
        targets = targets.to(device)

        outputs = my_nn(imgs)
        loss = loss_fn(outputs, targets)

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_step = train_step + 1
        if train_step % 100 == 0:
            print(f"训练次数：{train_step}，loss：{loss.item()}")
            writer.add_scalar("train_loss",loss.item(),train_step)

    test_total_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for imgs,targets in test_loader:
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = my_nn(imgs)
            loss = loss_fn(outputs,targets)
            test_total_loss = test_total_loss+loss.item()
            accuracy = (outputs.argmax(1)==targets).sum()
            total_accuracy = total_accuracy+accuracy

        print(f"第{i+1}轮测试，loss：{test_total_loss}")
        print(f"第{i + 1}轮测试准确率：{total_accuracy/len(test_data)}")
        writer.add_scalar("test_loss", test_total_loss, i+1)
        writer.add_scalar("test_accuracy", total_accuracy/len(test_data), i + 1)

        torch.save(my_nn,f"my_model_{i+1}.pth")
        print("模型已保存")

writer.close()

